import torch
import logging
import torchvision
import torch.nn.functional as F
from torch import nn
from model import vig
from model.layers import Flatten, L2Norm, GeM, Convs, NN


CHANNELS_NUM_IN_LAST_CONV = {
        "resnet18": 512,
        "resnet50": 2048,
        "resnet101": 2048,
        "resnet152": 2048,
        "vgg16": 512,
        "vig": 640
    }


def get_backbone(backbone_name):
    if backbone_name.startswith("resnet"):
        if backbone_name == "resnet18":
            backbone = torchvision.models.resnet18(pretrained=True)
        elif backbone_name == "resnet50":
            backbone = torchvision.models.resnet50(pretrained=True)
        elif backbone_name == "resnet101":
            backbone = torchvision.models.resnet101(pretrained=True)
        elif backbone_name == "resnet152":
            backbone = torchvision.models.resnet152(pretrained=True)

        for name, child in backbone.named_children():
            if name == "layer3":  # Freeze layers before conv_3
                break
            for params in child.parameters():
                params.requires_grad = False
        logging.debug(f"Train only layer3 and layer4 of the {backbone_name}, freeze the previous ones")
        layers = list(backbone.children())[:-2]  # Remove avg pooling and FC layer

    elif backbone_name == "vgg16":
        backbone = torchvision.models.vgg16(pretrained=True)
        layers = list(backbone.features.children())[:-2]  # Remove avg pooling and FC layer
        for l in layers[:-5]:
            for p in l.parameters(): p.requires_grad = False
        logging.debug("Train last layers of the VGG-16, freeze the previous ones")

    elif backbone_name == "vig":
        backbone = vig.pvig_s_224_gelu(pretrained=True)
        state_dict = torch.load('./pvig_s_82.1.pth.tar')
        backbone.load_state_dict(state_dict, strict=False)

        logging.debug("Train layers of the VIG, freeze no")

    if backbone_name != "vig":
        backbone = torch.nn.Sequential(*layers)

    features_dim = CHANNELS_NUM_IN_LAST_CONV[backbone_name]
    return backbone, features_dim


# Sensing direction module
class SDM(nn.Module):
    def __init__(self, out_c, in_c):
        super().__init__()
        self.conv = Convs(in_c, out_c)
        # self.line = NN()

    def forward(self, grd_coding, sat_coding):
        ori, option_ori = self.conv(grd_coding, sat_coding)
        ori = F.normalize(ori)
        option_ori = F.normalize(option_ori, dim=1)
        ori = torch.squeeze(ori, (0, 2, 3))
        option_ori = torch.squeeze(option_ori, 2)

        # ori, option_ori = self.line(ori, option_ori)
        # ori = F.normalize(ori, dim=0)
        # option_ori = F.normalize(option_ori, dim=1)

        return ori, option_ori


class GeoLocalizationNet(nn.Module):
    def __init__(self, backbone, fc_output_dim):
        super().__init__()
        self.backbone, features_dim = get_backbone(backbone)
        self.aggregation = nn.Sequential(
                L2Norm(),
                GeM(),
                Flatten(),
                nn.Linear(features_dim, fc_output_dim),
                L2Norm()
            )

    def forward(self, x):
        x = self.backbone(x)
        x = self.aggregation(x)
        return x


# class OriLocationNet(nn.Module):
#     def __init__(self, backbone):
#         super().__init__()
#         self.backbone, features_dim = get_backbone(backbone)
#         self.sdm = SDM
#
#     def forward(self, grd_x, sat_x):
#         grd_code = self.backbone(grd_x)
#         sat_code = self.backbone(sat_x)
#         ori, ori_list = self.sdm.forward(grd_code, sat_code)
#         return ori, ori_list
